import torch
import numpy as np
import glob
import os
import io
import random
import pickle
from torch.utils.data import Dataset, DataLoader
from lib.data.augmentation import Augmenter3D
from lib.utils.tools import read_pkl
from lib.utils.utils_data import flip_data
    
class MotionDataset(Dataset):
    def __init__(self, args, subset_list, data_split): # data_split: train/test
        np.random.seed(0)
        self.data_root = args.data_root
        self.subset_list = subset_list
        self.data_split = data_split
        file_list_all = []
        for subset in self.subset_list:
            data_path = os.path.join(self.data_root, subset, self.data_split)
            motion_list = sorted(os.listdir(data_path))
            for i in motion_list:
                file_list_all.append(os.path.join(data_path, i))
        self.file_list = file_list_all
        
    def __len__(self):
        'Denotes the total number of samples'
        return len(self.file_list)

    def __getitem__(self, index):
        raise NotImplementedError 

class MotionDataset3D(MotionDataset):
    def __init__(self, args, subset_list, data_split, noise_part=None):
        super(MotionDataset3D, self).__init__(args, subset_list, data_split)
        self.flip = args.flip
        self.synthetic = args.synthetic
        self.aug = Augmenter3D(args)
        self.gt_2d = args.gt_2d
        self.noise_part = noise_part

    def __getitem__(self, index):
        'Generates one sample of data'
        # Select sample
        file_path = self.file_list[index]
        motion_file = read_pkl(file_path)
        motion_3d = motion_file["data_label"]  
        if self.data_split=="train":
            if self.synthetic or self.gt_2d:
                motion_3d = self.aug.augment3D(motion_3d)
                motion_2d = np.zeros(motion_3d.shape, dtype=np.float32)
                motion_2d[:,:,:2] = motion_3d[:,:,:2]
                motion_2d[:,:,2] = 1                        # No 2D detection, use GT xy and c=1. # ？？？
            elif motion_file["data_input"] is not None:     # Have 2D detection 
                motion_2d = motion_file["data_input"]
                if self.flip and random.random() > 0.5:                        # Training augmentation - random flipping
                    motion_2d = flip_data(motion_2d)
                    motion_3d = flip_data(motion_3d)
            else:
                raise ValueError('Training illegal.') 
        elif self.data_split=="test":                                           
            motion_2d = motion_file["data_input"]
            if self.gt_2d:
                motion_2d[:,:,:2] = motion_3d[:,:,:2]
                motion_2d[:,:,2] = 1
            """ 对姿态数据做各个部分的噪声处理，模拟人体部位不可见情况下的识别结果 """
            motion_2d = add_part_noise(motion_2d, self.noise_part)
            # breakpoint()
        else:
            raise ValueError('Data split unknown.')    
        return torch.FloatTensor(motion_2d), torch.FloatTensor(motion_3d)
"""
In [7]: test_dataset = MotionDataset3D(args, args.subset_list, 'test', "upper")

In [11]: test_dataset[0]
> /home/hugoxana/Documents/MotionBERT/lib/data/dataset_motion_3d.py(72)__getitem__()
-> return torch.FloatTensor(motion_2d), torch.FloatTensor(motion_3d)
(Pdb) motion_2d.shape
(243, 17, 3)
(Pdb) motion_2d
array([[[-0.09399998, -0.19399998,  0.65554845],
        [-0.05400002, -0.19399998,  0.82664   ],
        [-0.03399998, -0.00399997,  0.77685416],
        ...,
        [-0.028     , -0.40199998,  0.78756124],
        [ 0.09000003, -0.36999997,  0.7411108 ],
        [ 0.16799998, -0.358     ,  0.8789877 ]],

       [[-0.09399998, -0.19399998,  0.648048  ],
        [-0.05400002, -0.19399998,  0.81554306],
        [-0.03399998, -0.00399997,  0.7651963 ],
        ...,
        [-0.028     , -0.40199998,  0.7935814 ],
        [ 0.08399999, -0.36999997,  0.7304168 ],
        [ 0.16799998, -0.358     ,  0.89171517]],

       [[-0.10000002, -0.19399998,  0.6493298 ],
        [-0.05400002, -0.19399998,  0.71486306],
        [-0.03399998, -0.00399997,  0.7662984 ],
        ...,
        [-0.028     , -0.40199998,  0.8040982 ],
        [ 0.09000003, -0.36999997,  0.77918804],
        [ 0.16799998, -0.358     ,  0.90233195]],

       ...,

       [[-0.32      , -0.19399998,  0.6618883 ],
        [-0.314     , -0.21199998,  0.7324008 ],
        [-0.3       , -0.02400002,  0.7489234 ],
        ...,
        [-0.30800003, -0.40800002,  0.6851989 ],
        [-0.288     , -0.298     ,  0.38696337],
        [-0.37800002, -0.27199998,  0.46720362]],

       [[-0.32      , -0.19399998,  0.6603777 ],
        [-0.30800003, -0.21199998,  0.75209576],
        [-0.302     , -0.02400002,  0.7647358 ],
        ...,
        [-0.302     , -0.40199998,  0.768924  ],
        [-0.282     , -0.29      ,  0.40990314],
        [-0.374     , -0.27800003,  0.48732424]],

       [[-0.328     , -0.19399998,  0.6582126 ],
        [-0.30800003, -0.21199998,  0.7876083 ],
        [-0.302     , -0.02400002,  0.6847276 ],
        ...,
        [-0.302     , -0.40199998,  0.755331  ],
        [-0.282     , -0.29      ,  0.44890186],
        [-0.374     , -0.284     ,  0.540006  ]]], dtype=float32)
"""
    
def add_part_noise(motion_2d, noise_part):
    """
    上半身：(7: 'belly', 8: 'neck', 9: 'nose', 10: 'head', 
			11: 'lsho', 12: 'lelb', 13: 'lwri', 
			14: 'rsho', 15: 'relb', 16: 'rwri')
    下半身：(0: 'root', 1: 'rhip', 2: 'rkne', 3: 'rank', 4: 'lhip', 5: 'lkne', 6: 'lank')
    左半身：(0: 'root', 1: 'rhip', 2: 'rkne', 3: 'rank', 7: 'belly', 8: 'neck', 9: 'nose', 10: 'head',
                    14: 'rsho', 15: 'relb', 16: 'rwri')
    右半身：(0: 'root', 4: 'lhip', 5: 'lkne', 6: 'lank', 7: 'belly', 8: 'neck', 9: 'nose', 10: 'head',
                    11: 'lsho', 12: 'lelb', 13: 'lwri')
    """
    if noise_part == "upper":
        noise_idx =  [7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
    elif noise_part == "lower":
        noise_idx = [0, 1, 2, 3, 4, 5, 6]
    elif noise_part == "left":
        noise_idx = [0, 1, 2, 3, 7, 8, 9, 10, 14, 15, 16]
    elif noise_part == "right":
        noise_idx = [0, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
    else:
        noise_part = None
        noise_idx = []
    # print("add noise to ", noise_part, " body part")
    """
    输入motion_2d的形状：(243, 17, 3)， 最后一个数字是置信度

    实际需要加入噪声的部分：[:, [部分的关键点索引], 0:2]
    这里尝试构造一个二维噪声，其长度为高斯分布，方向为均匀分布
    """
    std = 0.1
    mean = 0.1
    T = motion_2d.shape[0]
    J = motion_2d.shape[1]
    noise_len = np.random.randn(T, J) * std + mean
    noise_angle = np.random.rand(T, J) * 360 - 180
    noise_xy = np.zeros_like(motion_2d)
    for t in range(T):
        for j in range(J):
            if j in noise_idx:
                noise_xy[t][j][0] = noise_len[t, j] * np.cos(np.deg2rad(noise_angle[t, j]))
                noise_xy[t][j][1] = noise_len[t, j] * np.sin(np.deg2rad(noise_angle[t, j]))

    return motion_2d + noise_xy

